-module(websocket_server).
-compile(export_all).

-include("globals.hrl").
 
%% 开始执行
start() ->
    {ok, ListenSocket} = gen_tcp:listen(?PORT, [binary, {packet, 0}, {reuseaddr, true}, {active, true}]), %% 开始监听
    io:format("Starting websocket server, listening on port: ~p~n", [?PORT]),
    wait_connect(ListenSocket).

%% 等待连接
wait_connect(ListenSocket) ->
    {ok, Socket} = gen_tcp:accept(ListenSocket),
    Pid = spawn(?MODULE, get_request, [#session{socket = Socket}]),
    gen_tcp:controlling_process(Socket, Pid), %% 处理每个连接, 将处理过程加入新的进程
    wait_connect(ListenSocket).

%% 处理连接请求
get_request(Session) ->
    receive
        {tcp, Socket, HeaderData} ->
            HeaderList = binary:split(HeaderData, <<"\r\n">>, [global]),
            Header = [list_to_tuple(binary:split(Item, <<": ">>)) || Item <- HeaderList, Item /= nomatch],
            {_, SecWebSocketKey} = lists:keyfind(<<"Sec-WebSocket-Key">>, 1, Header),
            %% Sha1 = crypto:sha([SecWebSocketKey, <<"258EAFA5-E914-47DA-95CA-C5AB0DC85B11">>]),
            Sha1 = crypto:hash(sha, [SecWebSocketKey, <<"258EAFA5-E914-47DA-95CA-C5AB0DC85B11">>]),
            Base64 = base64:encode(Sha1),
            Handshake = [
                <<"HTTP/1.1 101 Switching Protocols\r\n">>,
                <<"Upgrade: websocket\r\n">>,
                <<"Connection: Upgrade\r\n">>,
                <<"Sec-WebSocket-Accept: ">>, Base64, <<"\r\n">>,
                <<"\r\n">>
            ],
            gen_tcp:send(Socket, Handshake), %% 发送给客户端进行验证
            loop(Session); %% 处理下一个连接请求
        _Any -> get_request(Session)
    end.

%% 循环处理数据
loop(Session) ->
    receive
        {tcp, _Socket, Data} -> 
            handle_data(Data, Session); %% 建立连接成功
        {tcp_closed, Socket} -> 
            gen_tcp:close(Socket);
        _Any -> 
            loop(Session)
    end.

%% 处理客户端提交的数据
handle_data(Data, Session) ->
    <<_Fin:1, _Rsv:3, _Opcode:4, _Mask:1, Len:7, Rest/binary>> = Data,
    <<Masking:4/binary, Payload:Len/binary, Next/binary>> = Rest,
    Line = unmask(Payload, Masking),
    case unicode:characters_to_list(Line) of
        {incomplete, _, _} ->
            gen_tcp:close(Session#session.socket);
        Str ->
            {NewSession, HandlerText} = websocket_handler:dispatcher(Str, Session), %% 处理经过 handler/model 之后的数据
            send_data(Session#session.socket, HandlerText), %% 发送数据到客户端
            case size(Next) of
                0 -> loop(NewSession); %% 如果发送完成, 则继续监听
                _Other -> handle_data(Next, Session) %% 否则, 继续进行下一次处理
            end
    end.

%% 发送数据给客户端
send_data(Socket, Data) -> 
    Bin = unicode:characters_to_binary(Data),
    Frame = <<1:1, 0:3, 1:4, 0:1, (size(Bin)):7, Bin/binary>>,
    gen_tcp:send(Socket, Frame).

%% 对客户端传上来的数据进行解码
unmask(Payload, Masking) -> unmask(Payload, Masking, <<>>).
unmask(Payload, Masking = <<MA:8, MB:8, MC:8, MD:8>>, Acc) ->
    case size(Payload) of
        0 -> Acc;
        1 -> 
            <<A:8>> = Payload,
            <<Acc/binary, (MA bxor A)>>;
        2 -> 
            <<A:8, B:8>> = Payload,
            <<Acc/binary, (MA bxor A), (MB bxor B)>>;
        3 -> 
            <<A:8, B:8, C:8>> = Payload,
            <<Acc/binary, (MA bxor A), (MB bxor B), (MC bxor C)>>;
        _Other ->
            <<A:8, B:8, C:8, D:8, Rest/binary>> = Payload,
            Acc1 = <<Acc/binary, (MA bxor A), (MB bxor B), (MC bxor C), (MD bxor D)>>,
            unmask(Rest, Masking, Acc1)
    end.